{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Streaming"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using LangGraph API\n",
    "\n",
    "Tell about server part of LangGraph Studio and prefered approach to build graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph_sdk import get_client\n",
    "\n",
    "URL = \"http://localhost:61693\"\n",
    "client = get_client(url=URL)\n",
    "\n",
    "# Search all hosted graphs\n",
    "assistants = await client.assistants.search()\n",
    "assistants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assistants[0][\"assistant_id\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "# Create a new thread\n",
    "thread = await client.threads.create()\n",
    "\n",
    "final_state = await client.runs.wait(\n",
    "    thread_id=thread[\"thread_id\"],\n",
    "    assistant_id=\"8a4ac7a4-50eb-5206-98cc-4a72345cb1f7\",\n",
    "    input={\"question\": \"Hi, I’m working on a Python project, and I’m stuck with handling API responses.\"}\n",
    ")\n",
    "\n",
    "final_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_state = await client.runs.wait(\n",
    "    thread_id=thread[\"thread_id\"],\n",
    "    assistant_id=\"8a4ac7a4-50eb-5206-98cc-4a72345cb1f7\",\n",
    "    input={\"question\": \"Sorry what was my previous question?\"}\n",
    ")\n",
    "\n",
    "final_state[\"answer\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_state = await client.runs.wait(\n",
    "    thread_id=thread[\"thread_id\"],\n",
    "    assistant_id=\"8a4ac7a4-50eb-5206-98cc-4a72345cb1f7\",\n",
    "    input={\"question\": \"Ahh, yeah right! So I’m mostly struggling with parsing JSON responses. Sometimes the structure isn’t what I expect, and it breaks my code.\"}\n",
    ")\n",
    "\n",
    "final_state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Streaming\n",
    "\n",
    "Observe the difference between constructing graph manually & using LangGraph Studio"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define chatbot graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from IPython.display import Image, display\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import MessagesState, StateGraph, START, END\n",
    "from langchain_core.messages import HumanMessage, SystemMessage, RemoveMessage\n",
    "\n",
    "\n",
    "# OPENAI_API_KEY environment variable must be set\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "# System message\n",
    "chatbot_system_message = SystemMessage(content=(\"\"\"\n",
    "You are a helpful and knowledgeable chatbot assistant. \n",
    "Your goal is to provide clear and accurate answers to user questions based on the information they provide. \n",
    "Stay focused, concise, and ensure your responses are relevant to the context of the conversation. \n",
    "If you don’t have enough information, ask for clarification.”\n",
    "\"\"\"))\n",
    "\n",
    "\n",
    "# Nodes\n",
    "def chatbot(state: MessagesState) -> MessagesState:\n",
    "    response = llm.invoke([chatbot_system_message] + state[\"messages\"]);\n",
    "    return MessagesState(messages = [response])\n",
    "\n",
    "\n",
    "# Graph\n",
    "workflow = StateGraph(MessagesState)\n",
    "workflow.add_node(chatbot)\n",
    "\n",
    "workflow.add_edge(START, \"chatbot\")\n",
    "workflow.add_edge(\"chatbot\", END)\n",
    "\n",
    "\n",
    "memory = MemorySaver()\n",
    "graph = workflow.compile(checkpointer=memory)\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Streaming modes:\n",
    "\n",
    "- updates (exposes only new data)\n",
    "- values (always shows the whole state)\n",
    "- messages\n",
    "- debug\n",
    "- custom"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stream_mode=updates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a thread\n",
    "config = {\"configurable\": {\"thread_id\": \"1\"}}\n",
    "\n",
    "user_input = HumanMessage(content=\"Hi, I’m working on a Python project, and I’m stuck with handling API responses.\")\n",
    "for event in graph.stream({\"messages\": [user_input]}, config, stream_mode=\"updates\"):\n",
    "    print(event)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_input = HumanMessage(content=\"Sorry what was my previous question?\")\n",
    "for event in graph.stream({\"messages\": [user_input]}, config, stream_mode=\"updates\"):\n",
    "    for m in event['chatbot']['messages']:\n",
    "        m.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stream_mode=values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"2\"}}\n",
    "\n",
    "user_input = HumanMessage(content=\"Hi, I’m working on a Python project, and I’m stuck with handling API responses.\")\n",
    "for event in graph.stream({\"messages\": [user_input]}, config, stream_mode=\"values\"):\n",
    "    print(event)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"2\"}}\n",
    "\n",
    "user_input = HumanMessage(content=\"Hi, I’m working on a Python project, and I’m stuck with handling API responses.\")\n",
    "for event in graph.stream({\"messages\": [user_input]}, config, stream_mode=\"values\"):\n",
    "    for m in event['messages']:\n",
    "        m.pretty_print()\n",
    "    print(\"\\n\")\n",
    "    print(\"#\"*100)\n",
    "    print(\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Streaming deeper (updates inside Node) - a.k.a. \"streaming LLM tokens from a specific node\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"4\"}}\n",
    "\n",
    "user_input = HumanMessage(content=\"Hi, I’m working on a Python project, and I’m stuck with handling API responses.\")\n",
    "for event in graph.stream({\"messages\": [user_input]}, config, stream_mode=\"messages\"):\n",
    "    print(event)\n",
    "\n",
    "# so we have a message with content and metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"4\"}}\n",
    "\n",
    "user_input = HumanMessage(content=\"Ahh, yeah right! So I’m mostly struggling with parsing JSON responses. Sometimes the structure isn’t what I expect, and it breaks my code.\")\n",
    "for msg, metadata in graph.stream({\"messages\": [user_input]}, config, stream_mode=\"messages\"):\n",
    "    if (metadata['langgraph_node'] == 'chatbot'):\n",
    "        print(msg.content, end=\"\")\n",
    "\n",
    "# same style of outputing data as in chat app (a token by token)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Streaming with LangGraph API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph_sdk import get_client\n",
    "\n",
    "URL = \"http://localhost:61693\"\n",
    "client = get_client(url=URL)\n",
    "\n",
    "assistants = await client.assistants.search()\n",
    "assistants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thread = await client.threads.create()\n",
    "\n",
    "input_message = HumanMessage(content=\"Hi, I’m working on a Python project, and I’m stuck with handling API responses.\")\n",
    "\n",
    "async for part in client.runs.stream(\n",
    "        thread[\"thread_id\"], \n",
    "        assistant_id=\"8a4ac7a4-50eb-5206-98cc-4a72345cb1f7\", \n",
    "        input={\"messages\": [input_message]}, \n",
    "        stream_mode=\"messages\"):\n",
    "    print(part)\n",
    "\n",
    "# check event types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import convert_to_messages\n",
    "\n",
    "thread = await client.threads.create()\n",
    "input_message = HumanMessage(content=\"Should I invest in Tesla stocks?\")\n",
    "\n",
    "async for event in client.runs.stream(\n",
    "            thread[\"thread_id\"], \n",
    "            assistant_id=\"b7480eb0-6390-53a5-9bc4-29bf27cbd1c4\", \n",
    "            input={\"messages\": [input_message]}, \n",
    "            stream_mode=\"values\"):\n",
    "    messages = event.data.get('messages',None)\n",
    "    if messages:\n",
    "        print(convert_to_messages(messages)[-1])\n",
    "\n",
    "# display content only with convert_to_messages util"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
